package cc.git.liuyan.customeraiagent.core.bigmodel.abstrct;

import cc.git.liuyan.customeraiagent.core.bigmodel.*;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

@Slf4j
public abstract class AbstractBigModel implements BigModel {
    private static final DateTimeFormatter dateTimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
    private final String name;
    protected final BigModelApiClient apiClient;
    protected BigModelChatTool chatTool;

    public AbstractBigModel(String name, BigModelApiClient apiClient) {
        this.name = name;
        this.apiClient = apiClient;
    }

    @Override
    public final String modelName() {
        return this.name;
    }

    @Override
    public final BigModelChatTool getChatTool() {
        return this.chatTool;
    }

    @Override
    public final void setChatTool(BigModelChatTool chatTool) {
        this.chatTool = chatTool;
    }

    @Override
    public Flux<BigModelChatOutputData> chatAsync(BigModelChatInputData inputData) {
        if (this.chatTool == null) {
            this.chatTool = new EmptyBigModelChatTool();
        }
        Flux<BigModelChatOutputData> flux = null;
        StringBuilder fluxAssistantResultBuilder = new StringBuilder();
        //自定义回复
        BigModelChatOutputData customReplay = this.chatTool.customReplay(inputData);
        if (customReplay != null) {
            flux = Flux.create(emitter -> {
                emitter.next(customReplay);
                emitter.complete();
            });
        } else {
            String prompt = inputData.getContent();
            //区分模型用途
            if (Objects.equals(inputData.getUsage(), "text2sql")) {
                prompt = this.chatTool.getText2SqlFinalPrompt(inputData);
            } else if (Objects.equals(inputData.getUsage(), "chat")) {
                prompt = this.chatTool.getChatFinalPrompt(inputData);
            } else {
                throw new IllegalArgumentException("Unsupported usage type: " + inputData.getUsage());
            }
            final String finalPrompt = prompt;

            flux = this.chatAsyncFinal(inputData, finalPrompt);
        }

        flux = flux.doOnSubscribe(subscription -> {
            log.info("Flux:doOnSubscribe");
            //多轮对话
            this.chatTool.saveChatData(inputData, new BigModelChatData(null, "user", inputData.getContent(), dateTimeFormatter.format(LocalDateTime.now())));
        }).doOnNext(outputData -> {
            //多轮对话
            if (!Objects.equals(outputData.getContent(), null)) {
                fluxAssistantResultBuilder.append(outputData.getContent());
            }
        }).doFinally(signalType -> {
            //多轮对话
            if (fluxAssistantResultBuilder.length() > 0) {
                this.chatTool.saveChatData(inputData, new BigModelChatData(null, "assistant", fluxAssistantResultBuilder.toString(), dateTimeFormatter.format(LocalDateTime.now())));
            }
            //用户问题持久化，可用于审查、安全分析
            this.chatTool.persistenceBigModelChatInputData(inputData, fluxAssistantResultBuilder.toString());
            log.info("Flux:doFinally");
        }).doOnError(throwable -> {
            log.error("Flux:onError", throwable);
            throw new RuntimeException(throwable);
        });

        return flux;
    }

    protected abstract Flux<BigModelChatOutputData> chatAsyncFinal(BigModelChatInputData inputData, String finalPrompt);
}
